'''
    由于主仓和私仓llama层名不一致, 私仓权重在主仓调用前需要转换
'''
from mindspore import load_checkpoint, save_checkpoint

ckpt_path = "/home/zhangsenzhen/2023Q2/checkpoint_download/Ziya-LLaMA-13B-v1.1-hf/ziya_old_output/rank_0/checkpoint_0.ckpt"
params = load_checkpoint(ckpt_path)

new_params = []
for k, v in params.items():
    if k == 'tok_embeddings.weight':
        new_k = 'model.tok_embeddings.embedding_weight'
    elif k == 'norm_out.weight':
        new_k = 'model.norm_out.weight'
    elif k == 'output.weight':
        new_k = 'lm_head.weight'
    else:
        new_k = 'model.layers.' + k
    
    new_params.append({"name": new_k, "data": v})

save_checkpoint(new_params, "/home/zhangsenzhen/2023Q2/checkpoint_download/llama/ziya/pinjie_output/rank_0/checkpoint_0.ckpt")
